import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LassoCV
from sklearn.feature_selection import mutual_info_regression
from sklearn.inspection import permutation_importance
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from xgboost import XGBRegressor
import shap
import time
import os
import matplotlib as mpl

# 设置绘图风格
sns.set_style("whitegrid")
plt.rcParams['figure.figsize'] = (12, 8)
plt.rcParams['font.size'] = 12
mpl.rcParams['font.family'] = 'SimHei'
mpl.rcParams['axes.unicode_minus'] = False


# 1. 加载数据（支持中文编码）
def load_data(filepath):
    try:
        # 尝试多种中文编码格式
        encodings = ['utf-8', 'gbk', 'gb2312', 'utf-8-sig', 'latin1']
        df = None
        used_encoding = None

        for encoding in encodings:
            try:
                df = pd.read_csv(filepath, parse_dates=['rectime'], encoding=encoding)
                used_encoding = encoding
                print(f"数据加载完成，使用编码: {encoding}")
                break
            except (UnicodeDecodeError, LookupError):
                continue
            except Exception as e:
                # 如果是日期解析错误，尝试不解析日期
                try:
                    df = pd.read_csv(filepath, encoding=encoding)
                    used_encoding = encoding
                    print(f"数据加载完成（跳过日期解析），使用编码: {encoding}")
                    break
                except:
                    continue

        if df is None:
            raise ValueError("无法用任何编码读取文件")

        print(f"数据维度: {df.shape[0]} 行, {df.shape[1]} 列")
        print(f"列名: {list(df.columns)}")
        return df
    except Exception as e:
        print(f"加载数据时出错: {e}")
        raise


# 2. 数据预处理（适配中文列名）
'''def preprocess_data(df, target_col):
    print("\n开始数据预处理...")

    # 检查目标列是否存在
    if target_col not in df.columns:
        available_cols = list(df.columns)
        print(f"目标列 '{target_col}' 不存在!")
        print(f"可用列名: {available_cols}")
        raise ValueError(f"目标列 '{target_col}' 不存在于数据中")

    # 处理缺失值
    numeric_cols = df.select_dtypes(include=[np.number]).columns
    if df[numeric_cols].isnull().sum().sum() > 0:
        null_count = df[numeric_cols].isnull().sum().sum()
        print(f"发现 {null_count} 个缺失值，使用中位数填充...")
        df[numeric_cols] = df[numeric_cols].fillna(df[numeric_cols].median())

    # 分离特征和目标
    # 排除非数值列和时间列
    exclude_cols = ['rectime', target_col]
    features = [col for col in df.columns if col not in exclude_cols and df[col].dtype in [np.number]]

    if not features:
        raise ValueError("没有找到数值型特征列")

    print(f"使用 {len(features)} 个数值特征进行分析")
    print(f"特征列表: {features}")

    X = df[features]
    y = df[target_col]

    # 标准化特征
    scaler = StandardScaler()
    X_scaled = pd.DataFrame(scaler.fit_transform(X), columns=X.columns)

    return X_scaled, y, features'''


def preprocess_data(df, target_col):
    print("\n开始数据预处理...")

    # 检查目标列是否存在
    if target_col not in df.columns:
        available_cols = list(df.columns)
        print(f"目标列 '{target_col}' 不存在!")
        print(f"可用列名: {available_cols}")
        raise ValueError(f"目标列 '{target_col}' 不存在于数据中")

    # 处理缺失值（只处理NaN，不处理0值）
    numeric_cols = df.select_dtypes(include=[np.number]).columns

    # 只计算真正的缺失值（NaN），不包括0
    null_count = df[numeric_cols].isnull().sum().sum()
    if null_count > 0:
        print(f"发现 {null_count} 个缺失值，使用中位数填充...")
        df[numeric_cols] = df[numeric_cols].fillna(df[numeric_cols].median())

    # 可选：统计0值的数量（但不处理）
    zero_counts = (df[numeric_cols] == 0).sum()
    total_zeros = zero_counts.sum()
    if total_zeros > 0:
        print(f"发现 {total_zeros} 个零值，将保留原值不作处理")
        # 可以打印每列零值的详细情况
        for col in zero_counts[zero_counts > 0].index:
            print(f"  列 '{col}': {zero_counts[col]} 个零值")

    # 分离特征和目标
    # 排除非数值列和时间列
    exclude_cols = ['rectime', target_col]
    features = [col for col in df.columns if col not in exclude_cols and df[col].dtype in [np.number]]

    if not features:
        raise ValueError("没有找到数值型特征列")

    print(f"使用 {len(features)} 个数值特征进行分析")
    print(f"特征列表: {features}")

    X = df[features]
    y = df[target_col]

    # 标准化特征
    scaler = StandardScaler()
    X_scaled = pd.DataFrame(scaler.fit_transform(X), columns=X.columns)

    return X_scaled, y, features


# 3. 计算多种相关性指标（保持不变）
def compute_feature_importance(X, y, features):
    print("\n开始计算特征重要性...")
    results = pd.DataFrame(index=features)

    # 1. 相关系数
    print("计算相关系数...")
    results['spearman_corr'] = X.corrwith(y, method='spearman')
    results['pearson_corr'] = X.corrwith(y, method='pearson')

    # 2. 互信息
    print("计算互信息...")
    mi = mutual_info_regression(X, y, random_state=42)
    results['mutual_info'] = mi

    # 3. 随机森林重要性
    print("训练随机森林模型...")
    rf = RandomForestRegressor(n_estimators=100, random_state=42, n_jobs=-1)
    rf.fit(X, y)
    results['rf_importance'] = rf.feature_importances_

    # 4. Lasso回归系数
    print("训练Lasso模型...")
    lasso = LassoCV(cv=5, random_state=42, n_jobs=-1)
    lasso.fit(X, y)
    results['lasso_coef'] = lasso.coef_

    # 5. XGBoost重要性
    print("训练XGBoost模型...")
    xgb = XGBRegressor(n_estimators=100, random_state=42, n_jobs=-1, verbosity=0)
    xgb.fit(X, y)
    results['xgb_importance'] = xgb.feature_importances_

    # 6. 排列重要性
    print("计算排列重要性...")
    perm_result = permutation_importance(rf, X, y, n_repeats=5, random_state=42, n_jobs=1)
    results['perm_importance'] = perm_result.importances_mean

    # 计算综合评分
    print("计算综合评分...")
    # 归一化所有指标到0-1范围
    normalized_results = results.abs().apply(lambda x: (x - x.min()) / (x.max() - x.min()))

    weights = {
        'spearman_corr': 0.10,
        'pearson_corr': 0.10,
        'mutual_info': 0.15,
        'rf_importance': 0.20,
        'lasso_coef': 0.15,
        'xgb_importance': 0.20,
        'perm_importance': 0.10
    }

    # 计算加权综合评分
    normalized_results['composite_score'] = 0
    for col, weight in weights.items():
        normalized_results['composite_score'] += normalized_results[col] * weight

    # 合并原始结果和综合评分
    final_results = pd.concat([results, normalized_results[['composite_score']]], axis=1)
    final_results = final_results.sort_values(by='composite_score', ascending=True)

    print("特征重要性计算完成!")
    return final_results


# 4. 可视化结果（适配中文显示）
def visualize_results(results, target_name):
    print("\n开始可视化结果...")

    # 创建输出目录
    os.makedirs(writefile_folder, exist_ok=True)

    # 1. 综合评分图
    plt.figure(figsize=(14, 10))
    results = results.sort_values(by='composite_score', ascending=True)

    plt.subplot(2, 1, 1)
    results['composite_score'].plot(kind='barh', color='skyblue')
    plt.title(f'综合特征相关性评分 (目标: {target_name})')
    plt.xlabel('综合评分 (0-1)')
    plt.ylabel('特征')
    plt.grid(axis='x', linestyle='--', alpha=0.7)

    # 设置y轴标签为中文特征名
    plt.yticks(rotation=0)

    # 2. 相关系数比较
    plt.subplot(2, 1, 2)
    corr_df = results[['spearman_corr', 'pearson_corr']].sort_values(by='spearman_corr', ascending=True)
    corr_df.plot(kind='barh', ax=plt.gca())
    plt.title('特征相关性分析')
    plt.xlabel('相关系数值')
    plt.ylabel('')
    plt.legend(['Spearman (非线性)', 'Pearson (线性)'])
    plt.grid(axis='x', linestyle='--', alpha=0.7)
    plt.yticks(rotation=0)

    plt.tight_layout()
    plt.savefig(os.path.join(writefile_folder, 'correlation_comparison.png'), dpi=300, bbox_inches='tight')
    # plt.show()

    # 3. 模型重要性比较
    plt.figure(figsize=(14, 8))
    model_df = results[['rf_importance', 'xgb_importance', 'perm_importance']].sort_values(by='rf_importance',
                                                                                           ascending=True)
    model_df.plot(kind='barh', ax=plt.gca())
    plt.title('模型特征重要性比较')
    plt.xlabel('标准化重要性值')
    plt.ylabel('特征')
    plt.legend(['随机森林', 'XGBoost', '排列重要性'])
    plt.grid(axis='x', linestyle='--', alpha=0.7)
    plt.yticks(rotation=0)

    plt.tight_layout()
    plt.savefig(os.path.join(writefile_folder, 'model_importance_comparison.png'), dpi=300, bbox_inches='tight')
    # plt.show()

    # 4. 综合评分热力图
    plt.figure(figsize=(14, 10))
    corr_cols = ['spearman_corr', 'pearson_corr', 'mutual_info',
                 'rf_importance', 'xgb_importance', 'lasso_coef',
                 'perm_importance', 'composite_score']

    # 按综合评分排序
    sorted_results = results.sort_values(by='composite_score', ascending=False)

    # 创建相关性矩阵
    corr_matrix = sorted_results[corr_cols].corr()

    # 绘制热力图
    sns.heatmap(corr_matrix, annot=True, fmt=".2f", cmap="coolwarm",
                cbar_kws={'label': '相关性系数'}, annot_kws={"size": 10})
    plt.title('特征评估指标相关性分析')
    plt.xticks(rotation=45, ha='right')
    plt.tight_layout()
    plt.savefig(os.path.join(writefile_folder, 'metric_correlation_heatmap.png'), dpi=300, bbox_inches='tight')
    # plt.show()


# 配置参数
FILE_PATH = r"F:\dataset\BQW_min\data_min\A1_12month.csv"  # 替换文件路径
TARGET_COL = '平均齿轮箱冷却水温度'  # 替换为目标列名

print("=" * 60)
print("特征相关性分析系统")
print("=" * 60)

start_time = time.time()
try:
    # 加载数据
    df = load_data(FILE_PATH)

    # 可选：删除不需要的列（如果需要）
    '''df = df.drop(columns=['平均发电机绕组温度1', '平均发电机绕组温度2','平均发电机绕组温度4',
                          '平均发电机绕组温度5','平均发电机绕组温度6'])'''
    # df = df.drop(columns=['平均发电机绕组温度2', '平均发电机绕组温度1', '平均发电机绕组温度4','平均发电机绕组温度3','平均发电机绕组温度5'])
    # df = df.drop(columns=['平均齿轮箱水泵出口温度','平均齿轮箱水泵1入口温度','平均齿轮箱水泵2入口温度','平均齿轮箱冷却器出口油温'])
    df = df.drop(columns=['平均发电机内冷入口温度1','平均发电机内冷入口温度2','平均发电机内冷出口温度1','平均发电机内冷出口温度2',
                          '平均发电机外冷入口温度1','平均发电机外冷出口温度2','平均发电机外冷出口温度1','grgeneratoroutsiderecyclewindoutlettemperature2_avg',
                          '平均齿轮箱水泵出口温度'
                       ])
    # 数据预处理
    X, y, features = preprocess_data(df, TARGET_COL)

    # 计算特征重要性
    importance_df = compute_feature_importance(X, y, features)

    # 保存结果
    writefile_folder = r"F:\time_series_my_algorithm\60min\feature_importance\grGeneratorNacelleSideBearingTemperature\new2"
    trainfile_name = "grGeneratorNacelleSideBearingTemperature_importance.csv"

    # 自动创建目录（如果不存在）
    os.makedirs(writefile_folder, exist_ok=True)
    trainfile_path = os.path.join(writefile_folder, trainfile_name)
    importance_df.to_csv(trainfile_path, encoding='utf-8-sig')  # 使用utf-8-sig保存中文
    print(f"\n结果已保存至 '{trainfile_path}'")

    # 可视化
    visualize_results(importance_df, TARGET_COL)

    # 打印最重要的特征
    print("\n" + "=" * 60)
    print("特征相关性排名 (从低到高):")
    print("=" * 60)
    print(importance_df[['composite_score']].sort_values(by='composite_score', ascending=True))

    # 打印最重要的5个特征
    print("\n" + "=" * 60)
    print("最重要的5个特征:")
    print("=" * 60)
    top_features = importance_df.sort_values(by='composite_score', ascending=False).head(5)
    print(top_features[['composite_score']])

    print(f"\n分析完成! 总耗时: {time.time() - start_time:.2f}秒")

except Exception as e:
    print(f"\n程序执行出错: {e}")
    import traceback

    traceback.print_exc()
